#!/usr/bin/python

from SimpleXMLRPCServer import SimpleXMLRPCServer
import socket
import sys
import pwd
import os
import time
import atexit
from signal import SIGTERM
import tempfile
import subprocess
import md5
import fileinput

try:
    import guestfs
except ImportError:
    print "Could not find python-libguestfs, Please make sure it is installed."
    sys.exit(1)


IPTABLES_TOOL = '/sbin/iptables'
SERVICE_TOOL = "/sbin/service"
LVS_TOOL = 'lvs'
PGREP_TOOL = '/usr/bin/pgrep'
PVCREATE = 'pvcreate'
VGCREATE = 'vgcreate'
VGREMOVE = 'vgremove'
LVM = '/sbin/lvm'
DMSETUP = '/sbin/dmsetup'
MV = '/bin/mv'
RM = '/bin/rm'
CP = '/bin/cp'

QEMU_PATH = '/usr/libexec/qemu-kvm'
VDSM_PID_FILE = '/var/vdsm/vdsmd.pid'
LVM_CONF = '/etc/lvm/lvm.conf'
LVM_CONF_BACKUP = '.backup'
LVM_CACHE = '/etc/lvm/cache/.cache'

NFS_PORT = 2049

# Volume Format
COW_FORMAT = 4
RAW_FORMAT = 5

IMAGE_FORMAT_TRANSLATE = {COW_FORMAT : "qcow2", RAW_FORMAT : "raw"}

class Daemon(object):
    """
    A generic daemon class.

    Usage: subclass the Daemon class and override the run() method
    """
    def __init__(self, pidfile=None, stdin='/dev/null', stdout='/dev/null', stderr='/dev/null'):
        if pidfile is None:
            pidfile = '/tmp/%s.pid' % self.__class__.__name__

        self.stdin = stdin
        self.stdout = stdout
        self.stderr = stderr
        self.pidfile = pidfile

    def readPidFile(self):
        """
        Reads the PID file and makes sure it is valid.

        :returns: the pid of the already running daemon or :keyword:`None` if there is no PID file
                  or the process it refers to doesn't exist.
        """
        try:
            pidFile = file(self.pidfile,'r')
            pid = int(pidFile.read().strip())
            pidFile.close()
        except IOError:
            pid = None

        if not pid is None:
            try:
                os.kill(pid, 0)
            except OSError:
                sys.stderr.write("Found a pid file but process '%d' does not exist. Deleteing old pid file.\n" % pid)
                self.delpid()
                pid = None

        return pid

    def daemonize(self):
        """
        do the UNIX double-fork magic
        """
        try:
            pid = os.fork()
            if pid > 0:
                # exit first parent
                while (self.readPidFile() is None):
                    time.sleep(0.2)
                sys.exit(0)
        except OSError, ex:
            sys.stderr.write("fork #1 failed: %d (%s)\n" % (ex.errno, ex.strerror))
            sys.exit(1)

        # decouple from parent environment
        os.chdir("/")
        os.setsid()
        os.umask(0)

        # do second fork
        try:
            pid = os.fork()
            if pid > 0:
                # exit from second parent
                sys.exit(0)
        except OSError, ex:
            sys.stderr.write("fork #2 failed: %d (%s)\n" % (ex.errno, ex.strerror))
            sys.exit(1)

        # redirect standard file descriptors
        sys.stdout.flush()
        sys.stderr.flush()
        stdIn = file(self.stdin, 'r')
        stdOut = file(self.stdout, 'a+')
        stdErr = file(self.stderr, 'a+', 0)
        os.dup2(stdIn.fileno(), sys.stdin.fileno())
        os.dup2(stdOut.fileno(), sys.stdout.fileno())
        os.dup2(stdErr.fileno(), sys.stderr.fileno())

        # write pidfile
        atexit.register(self.delpid)
        pid = str(os.getpid())
        file(self.pidfile,'w+').write("%s\n" % pid)

    def delpid(self):
        os.remove(self.pidfile)

    def status(self):
        """
        Find out if the daemon is running
        """

        return ((self.readPidFile() is None) == False)

    def start(self):
        """
        Start the daemon
        """
        print "Starting Ruth Agent..."
        # Check for a pidfile to see if the daemon already runs
        if not self.readPidFile() is None:
            message = "pidfile %s already exist. Daemon already running?\n"
            sys.stderr.write(message % self.pidfile)
            sys.exit(1)

        # Start the daemon
        self.daemonize()
        self.run()

    def stop(self):
        """
        Stop the daemon
        """
        print "Stopping Ruth Agent..."
        # Get the pid from the pidfile
        pid = self.readPidFile()

        if pid is None:
            message = "pidfile %s does not exist. Daemon not running?\n"
            sys.stderr.write(message % self.pidfile)
            return # not an error in a restart

        # Try killing the daemon process
        try:
            while 1:
                os.kill(pid, SIGTERM)
                time.sleep(0.1)
        except OSError, err:
            err = str(err)
            if err.find("No such process") > 0:
                if os.path.exists(self.pidfile):
                    os.remove(self.pidfile)
            else:
                print str(err)
                sys.exit(1)

        print "Ruth Agent Stopped."

    def restart(self):
        """
        Restart the daemon
        """
        self.stop()
        self.start()

    def run(self):
        """
        You should override this method when you subclass Daemon. It will be called after the process has been
        daemonized by start() or restart().
        """


FAKE_INTERPRETER_SOCKET_TEMPLATE = "/var/vdsm/%s.interpreter.socket"

class ImpersonationStack (object):
    def __init__(self):
        self._stack = []

    def push(self, uid):
        self._stack.append(os.geteuid())
        try:
            os.seteuid(uid)
        except:
            self.pop()
            raise

    def pop(self):
        os.seteuid(self._stack.pop())

    def __del__(self):
        """
        Makes sure we are back where we started when the object dies.
        """
        while len(self._stack) > 0:
            self.pop()

class RuthAgent(Daemon):
    """
    An agent that sits on the server where a VDSM daemon is running and helps with testing
    """
    def run(self):
        server = SimpleXMLRPCServer(("0.0.0.0", 3099))
        server.register_multicall_functions()
        server.register_instance(self)
        server.serve_forever()

    def _getVdsmUid(self):
        """
        Get the UID of the VDSM user.
        """
        return pwd.getpwnam('vdsm').pw_uid

    def formatDiskImage(self, drivePath, partitions, format=None):
        """
        Format a disk image.

        :param drivePath: The path to the image file on the host.
        :param partitions: A list of tuples ``(size, fstype)``
        """
        # Because I want to do as little as possible as VDSM I prepare everything beforehand

        # Prepare signature file. It will be used to check for file system integrity.
        srcFd, srcPath = tempfile.mkstemp()
        tempFile = os.fdopen(srcFd, "wb")
        timeStamp = str(time.time())
        checksum = md5.new(timeStamp).hexdigest()

        tempFile.write(timeStamp + ":" + checksum)
        tempFile.flush()
        tempFile.close()

        os.chmod(srcPath, 0666)

        # Prepare the parameters for the disk formatter (see `man sfdisk`)
        fdiskLines = ["0,%d" % partitions[0][0]]
        for size, fsType in partitions[1:]:
            fdiskLines.append(",%d" % size)

        #Start impersonation
        impStack = ImpersonationStack()
        impStack.push(self._getVdsmUid())
        try:
            guest = guestfs.GuestFS()
            guest.set_qemu(QEMU_PATH)

            if format is None:
                guest.add_drive(drivePath)
            else:
                if not os.path.lexists(drivePath):
                    raise OSError("File Not found")

                guest.config('-drive', "file=%s,format=%s" % (drivePath, IMAGE_FORMAT_TRANSLATE[format]))

            guest.launch()

            device = guest.list_devices()[0]
            guest.sfdisk(device, 0, 0, 0, fdiskLines)

            # Format partitions
            for i in range(len(partitions)):
                fsType = partitions[i][1]
                guest.mkfs(fsType, "%s%d" % (device, (i + 1)))

            guest.mount(device + '1', "/")

            guest.upload(srcPath, "/.sigfile")

            guest.sync()
            guest.umount_all()
        finally:
            impStack.pop()

        os.unlink(srcPath)

        return True

    def checkDiskImageSignature(self, drivePath, format=None):
        srcFd, srcPath = tempfile.mkstemp()
        os.unlink(srcPath)

        impStack = ImpersonationStack()
        impStack.push(self._getVdsmUid())
        try:
            try:
                guest = guestfs.GuestFS ()
                guest.set_qemu(QEMU_PATH)

                if format is None:
                    guest.add_drive(drivePath)
                else:
                    if not os.path.lexists(drivePath):
                        raise OSError("File Not found")

                    guest.config('-drive', "file=%s,format=%s" % (drivePath, IMAGE_FORMAT_TRANSLATE[format]))

                guest.launch()

                device = guest.list_devices()[0]
                guest.mount(device + '1', "/")

                guest.download("/.sigfile", srcPath)

                guest.sync()
                guest.umount_all()
            except:
                return False
        finally:
            impStack.pop()

        tempFile = open(srcPath, "rb")

        timestamp, checksum = tempFile.read().split(":")

        tempFile.close()
        os.unlink(srcPath)

        return (md5.new(timestamp).hexdigest() == checksum)

    def addFileToDiskImage(self, drivePath, partitionNumber, targetPath, content, format=None):
        """
        Adds a file to a file system on a disk image.

        :param drivePath: The path to the image file on the host.
        :param partitionNumber: The partition number on the image disk to copy the file to.
        :param targetPath: The path relative to the root of the filesystem on the
                           selected partition where you want the file to be created.
        :param content: The content of the file.
        """
        srcFd, srcPath = tempfile.mkstemp()
        tempFile = os.fdopen(srcFd, "wb")
        tempFile.write(content)
        tempFile.flush()
        tempFile.close()

        os.chmod(srcPath, 0666)

        impStack = ImpersonationStack()
        impStack.push(self._getVdsmUid())
        try:
            guest = guestfs.GuestFS ()
            guest.set_qemu(QEMU_PATH)

            if format is None:
                guest.add_drive(drivePath)
            else:
                if not os.path.lexists(drivePath):
                    raise OSError("File Not found")

                guest.config('-drive', "file=%s,format=%s" % (drivePath, IMAGE_FORMAT_TRANSLATE[format]))

            guest.launch()

            device = guest.list_devices()[0]
            guest.mount("%s%d" % (device, partitionNumber), "/")

            guest.upload(srcPath, targetPath)

            guest.sync()
            guest.umount_all()
        finally:
            impStack.pop()

        os.unlink(srcPath)

        return True

    def _connectToFakeVM(self, vmUUID):
        socketAddress = FAKE_INTERPRETER_SOCKET_TEMPLATE % vmUUID
        sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        sock.connect(socketAddress)
        return sock

    def ping(self):
        """
        Return True

        For connection testing purposes only
        """
        return True

    def checkVolumeAvailability(self, volumes):
        """
        Checks if the lvs representing the specified volumes are available.

        :param volumes: a list of tuples ``(sdUUID, volUUID)``

        :returns: a list containing matching results (True, False, None...)
                  None will be used if lv doesn't exist.
        """
        cmd = [LVM, LVS_TOOL, '-o', 'lv_name,vg_name,attr']

        for vgName, lvName in volumes:
            cmd.append("%s/%s" % (vgName, lvName))

        out, err = self._exec(cmd)

        lines = out.splitlines()[1:]
        res = [None] * len(volumes)
        for line in lines:
            lv, vg, state = line.strip().split()
            isAvailable = (state[4] == 'a')

            for i in range(len(volumes)):
                vgName, lvName = volumes[i]

                if vgName == vg and lvName == lv:
                    res[i] = isAvailable

        return res

    def crashVDSM(self):
        otherPID = open(VDSM_PID_FILE).read().strip()
        out, err = self._exec([PGREP_TOOL, "vdsm"])

        vdsmPID = None
        out = out.splitlines()
        for pid in out:
            if int(pid) != int(otherPID):
                vdsmPID = int(pid)
                break

        if vdsmPID == None:
            raise RuntimeError("Could not find out who is VDSM. Is it running?")

        os.kill(vdsmPID, 9)
        return True

    def resetFakeVMInterpreted(self, vmUUID):
        self.sendCommandToFakeVM(vmUUID, "#RESET")

    def sendCommandToFakeVM(self, vmUUID, command):
        """
        Send a command to be executed by fakeQemu.
        """
        sock = self._connectToFakeVM(vmUUID)

        if not command.endswith("\n"):
            command += "\n"

        sock.send(command)
        return True

    def readFromFakeVMInterpreter(self, vmUUID):
        """
        Read a result from the interpreter
        """
        sock = self._connectToFakeVM(vmUUID)

        sock.settimeout(1.0)
        try:
            size = int(sock.recv(4), 16)
        except socket.timeout:
            return ""

        sock.settimeout(None)

        if size == 0:
            return ""
        return sock.recv(size)

    def makeFakeVMRaiseEvent(self, vmUUID, msg):
        """
        Push a message to a fakeqemu event queue.
        """
        command = "events.append('%s')\n" % msg
        self.sendCommandToFakeVM(vmUUID, command)
        return True

    def makeFakeVMRaiseWatermarkReached(self, vmUUID, blockDevice, allocated=799, watermark=798):
        """
        Make a fakeQemu send a 'high watermark reached' to the vdsm monitor.
        """
        msg = "# high watermark reached for %s alloc=%d mark=%d" % (blockDevice, allocated, watermark)
        self.makeFakeVMRaiseEvent(vmUUID, msg)
        return True

    def blockConnectionToHost(self, connection):
        """
        Block network to a storage connection using iptables.
        """
        name, conType, con = connection
        if conType == 1: #NFS
            self._exec([IPTABLES_TOOL, '-A', 'OUTPUT', '-d', con['connection'].split(":")[0], '-p', 'tcp', '--dport', str(NFS_PORT), '-j', 'DROP'])
        else:
            self._exec([IPTABLES_TOOL, '-A', 'OUTPUT', '-d', con['connection']]) #TODO : Block specific iscsi port, Handle FCP somehow

        return True

    def unblockConnectionToHost(self, connection):
        """
        Unblocks a previously blocked connection.
        """
        name, conType, con = connection
        if conType == 1: #NFS
            self._exec([IPTABLES_TOOL, '-D', 'OUTPUT', '-d', con['connection'].split(":")[0], '-p', 'tcp', '--dport', str(NFS_PORT), '-j', 'DROP'])
        else:
            self._exec([IPTABLES_TOOL, '-D', 'OUTPUT', '-d', con['connection']]) #TODO : Block specific iscsi port, Handle FCP somehow

        return True

    def flushConnectionBlock(self):
        """
        Clear all connection blocks made.
        """
        self._exec([IPTABLES_TOOL, "--flush"])

    def activateLV(self, vg, lv):
        """
        Activate a given LV (vg/lv)

        :param vg: VG name
        :param lv: LV name
        """
        self._exec([LVM, 'lvchange', '-a', 'y', "%s/%s" % (vg, lv)])

        return True

    def removeFile(self, path):
        """
        Remove the given file (path)

        :param path: File name to remove
        """
        if os.path.lexists(path):
            os.unlink(path)

        return True

    def restartVdsmService(self):
        """
        Restart VDSM.
        """
        self._exec([SERVICE_TOOL, "vdsmd", "restart"])
        return True

    def getMPDevList(self):
        """
        Return list of all multipath devices configured on the host
        """
        out = self._exec([DMSETUP, "ls", "--target=multipath"])[0].split("\n")
        return [line.strip() for line in out if len(line.strip()) > 0]

    def getLVMDevList(self, useVDSMFilter=False):
        """
        Return list of devices that lvm "sees"
        """
        env = None
        if useVDSMFilter:
            env = {"LVM_SYSTEM_DIR": "/var/vdsm/lvm"}
        cmd = [LVM, "lvmdiskscan"]
        out = self._exec(cmd, env)[0].split("\n")
        return [line.strip() for line in out if len(line.strip()) > 0]

    def createVG(self, vgName, devlist):
        """
        Create a VG
        """
        devlist = [os.path.join("/dev/mapper/", dev) for dev in devlist]
        for device in devlist:
            self._exec([LVM, PVCREATE, "-ff", device])

        cmd = [LVM, VGCREATE, vgName]
        cmd.extend(devlist)
        self._exec(cmd)

        return vgName

    def removeVG(self, vgName):
        """
        Force remove a VG
        """
        cmd = [LVM, VGREMOVE, "-f", vgName]
        self._exec(cmd)
        return vgName

    def changeLvmConf(self, restore=False):
        if restore:
            # restore original lvm.conf
            cmd = [MV, "-f", LVM_CONF + LVM_CONF_BACKUP, LVM_CONF]
            self._exec(cmd)
        else:
            for line in fileinput.FileInput(LVM_CONF, inplace=1, backup=LVM_CONF_BACKUP):
                if "preferred_names" in line:
                    line=line.replace('/dev/mapper', '/dev/mpath')
                sys.stdout.write(line)
        return True

    def rebuildLvmCache(self):
        # Remove prev. cache
        if os.path.exists(LVM_CACHE):
            os.unlink(LVM_CACHE)
        # Rebuild new cache
        cmd = [LVM, LVS_TOOL]
        self._exec(cmd)
        return True

    def _transformPath(self, remotePath):
        """
        Transform remote path to new one for local mount
        """
        return remotePath.replace('_','__').replace('/','_')

    def duplicateDomain(self, oldCon, newCon):
        prefix = '/rhev/data-center/mnt'
        src = self._transformPath(oldCon)
        dst = self._transformPath(newCon)

        # Duplicate domain
        cmd = ['su', 'vdsm', '-c', ' '.join([CP, "-R", os.path.join(prefix, src, '*'), os.path.join(prefix, dst)])]
        self._exec(cmd)
        return True

    def eraseDomain(self, con):
        prefix = '/rhev/data-center/mnt'
        src = self._transformPath(con)

        # Erase domain
        cmd = ['su', 'vdsm', '-c', ' '.join([RM, "-rf", os.path.join(prefix, src, '*')])]
        self._exec(cmd)
        return True

    @classmethod
    def _exec(cls, cmd, env=None):
        process = subprocess.Popen(cmd, env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        out, err = process.communicate()
        rc = process.returncode
        if rc != 0:
            raise Exception("exec command failed (CMD:%s) (OUT:%s) (ERR:%s)" % (subprocess.list2cmdline(cmd), out, err))

        return out, err

def main():
    daemon = RuthAgent()
    if len(sys.argv) == 2:
        if 'start' == sys.argv[1]:
            daemon.start()
        elif 'stop' == sys.argv[1]:
            daemon.stop()
        elif 'restart' == sys.argv[1]:
            daemon.restart()
        elif 'status' == sys.argv[1]:
            if daemon.status():
                print "Running"
            else:
                print "Not running"
        else:
            print "Unknown command"
            sys.exit(2)
        sys.exit(0)
    else:
        print "usage: %s start|stop|restart" % sys.argv[0]
        sys.exit(2)
    sys.exit(0)

if __name__ == "__main__":
    main()
